package cn.lijiajia3515.cairo.auth.config.security;

import cn.hutool.core.codec.Base64Decoder;
import cn.lijiajia3515.cairo.auth.framework.security.oauth2.OAuth2JwkProperties;
import cn.lijiajia3515.cairo.auth.framework.security.oauth2.resourceserver.jwt.jose.Jwks;
import cn.lijiajia3515.cairo.core.CoreConstants;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.util.Assert;

import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.ArrayList;
import java.util.List;

/**
 * jwt sets config
 */
@Configuration(proxyBeanMethods = false)
@AutoConfigureAfter(OAuthServerBeanConfig.class)
public class JwtSetConfig {

	/**
	 * 配置文件配置 jwk
	 */

	@Bean
	// @ConditionalOnBean(OAuth2JwkProperties.class)
	@Primary
	@Order(Ordered.HIGHEST_PRECEDENCE)
	JWKSource<SecurityContext> cairoJwkSource(OAuth2JwkProperties jwk) {
		List<JWK> keys = new ArrayList<>();
		jwk.getKeys().forEach((kid, key) -> {
			String privateKeyStr = key.getPrivateKey();
			String publicKeyStr = key.getPublicKey();
			Assert.notNull(privateKeyStr, "jwk private key 不能设置为空");
			Assert.notNull(publicKeyStr, "jwk public key 不能设置为空");
			PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(Base64Decoder.decode(privateKeyStr.getBytes()));
			X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(Base64Decoder.decode(publicKeyStr.getBytes()));
			try {
				RSAPrivateKey privateKey = (RSAPrivateKey) KeyFactory.getInstance("RSA").generatePrivate(privateKeySpec);
				RSAPublicKey publicKey = (RSAPublicKey) KeyFactory.getInstance("RSA").generatePublic(publicKeySpec);
				keys.add(new RSAKey.Builder(publicKey).privateKey(privateKey).keyID(kid).build());
			} catch (InvalidKeySpecException | NoSuchAlgorithmException e) {
				throw new IllegalStateException("配置错误", e);
			}
		});
		JWKSet jwkSet = new JWKSet(keys);
		return (jwkSelector, securityContext) -> jwkSelector.select(jwkSet);
	}

	@Bean
	@ConditionalOnMissingBean
	@Order(Ordered.LOWEST_PRECEDENCE)
	JWKSource<SecurityContext> defaultJwkSource() {
		JWK jwk = Jwks.generateRsa(CoreConstants.SNOWFLAKE.nextIdStr());
		JWKSet jwkSet = new JWKSet(jwk);
		return (jwkSelector, securityContext) -> jwkSelector.select(jwkSet);
	}

}
